/*
 * Copyright 2010 Kevin L'Huillier <klhuillier@gmail.com>
 * 
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not
 * use this file except in compliance with the License. You may obtain a copy of
 * the License at
 * 
 * http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 * License for the specific language governing permissions and limitations under
 * the License.
 */

package org.lhuillier.pwsafe.io.v3;


import java.io.ByteArrayOutputStream;
import java.security.SecureRandom;
import java.util.Arrays;
import org.lhuillier.pwsafe.io.RawDb;
import org.lhuillier.pwsafe.io.StretchedKey;
import org.lhuillier.pwsafe.io.codec.BinConverter;
import org.lhuillier.pwsafe.io.codec.RawRecord;
import org.lhuillier.pwsafe.model.RawField;
import com.google.inject.Inject;


class EncryptImpl implements Encrypt {
    private final SecureRandom random = new SecureRandom();
    private final BinConverter binConv;
    private final Crypto crypto;

    @Inject
    EncryptImpl(BinConverter binConv, Crypto crypto) {
        this.binConv = binConv;
        this.crypto = crypto;
    }

    @Override
    public byte[] encrypt(RawDb db, StretchedKey stretchedKey) {
        byte[] salt = stretchedKey.getSalt();
        byte[] key = stretchedKey.getKey();
        byte[] recordKey = generateRandom(32);
        byte[] hmacKey = generateRandom(32);
        byte[] serialized = serializeData(db);
        byte[] hmac = crypto.calculateHmac(serialized, hmacKey);
        byte[] iv = generateRandom(16);
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        out.write(new byte[] {'P', 'W', 'S', '3'}, 0, 4);
        out.write(salt, 0, 32);
        out.write(binConv.writeInt(stretchedKey.getIterations()), 0, 4);
        out.write(crypto.digestKey(key), 0, 32);
        out.write(crypto.encryptKey(recordKey, key), 0, 32);
        out.write(crypto.encryptKey(hmacKey, key), 0, 32);
        out.write(iv, 0, 16);
        byte[] encryptedData = crypto.encryptRecords(serialized, recordKey, iv);
        out.write(encryptedData, 0, encryptedData.length);
        out.write(binConv.writeString("PWS3-EOFPWS3-EOF"), 0, 16);
        out.write(hmac, 0, 32);
        return out.toByteArray();
    }
    
    /**
     * @return a new array with the specified number of random bytes
     */
    private byte[] generateRandom(int size) {
        byte[] result = new byte[size];
        random.nextBytes(result);
        return result;
    }

    private byte[] serializeData(RawDb db) {
        ByteArrayOutputStream out = new ByteArrayOutputStream();

        for (RawField header : db.getHeaders()) {
            byte[] data = pad(toBlocks(header));
            out.write(data, 0, data.length);
        }

        // Do not reuse eof as it should contain 11 random bytes.
        out.write(eof(), 0, 16);

        for (RawRecord record : db.getRecords()) {
            for (RawField field : record.getFields()) {
                byte[] data = pad(toBlocks(field));
                out.write(data, 0, data.length);
            }
            out.write(eof(), 0, 16);
        }

        return out.toByteArray();
    }

    private byte[] pad(byte[] blocks) {
        int reqd = 16 - (blocks.length % 16);
        if (reqd == 16) {
            return blocks;
        }
        byte[] result = Arrays.copyOf(blocks, blocks.length + reqd);
        byte[] rand = generateRandom(reqd);
        System.arraycopy(rand, 0, result, blocks.length, reqd);
        return result;
    }

    private byte[] toBlocks(RawField field) {
        int len = field.getValue().length;
        byte[] result = new byte[len + 5];
        result[0] = (byte) (len & 0xff);
        result[1] = (byte) (len >> 8 & 0xff);
        result[2] = (byte) (len >> 16 & 0xff);
        result[3] = (byte) (len >> 24 & 0xff);
        result[4] = field.getTypeId();
        System.arraycopy(field.getValue(), 0, result, 5, len);
        return result;
    }

    private byte[] eof() {
        byte[] result = new byte[16];
        random.nextBytes(result);

        // Set first 4 bytes to 0 to indicate the data length is 0
        for (int i = 0; i < 4; i++) {
            result[i] = 0;
        }

        // Add the type identifier. It is 0xff for both the header and records
        result[4] = (byte) 0xff;
        return result;
    }
}
